import numpy as np
from typing_extensions import Iterator
from ...base import DatabaseBuilder


def latlon2pix(lat, lon, radius, depth):
    """
    将经纬度转换为球面坐标，注意单位是弧度
    """
    x = (radius + depth) * np.cos(lat) * np.cos(lon)
    y = (radius + depth) * np.cos(lat) * np.sin(lon)
    z = (radius + depth) * np.sin(lat)
    return x, y, z


class BinadHierarchyBuilder(DatabaseBuilder):
    """
    基于二元组的不变量库构建基类

    使用普通方式构建的导航数据库，即不作任何删减，只将原始数据按视场大小划分
    划分方式为最高点处反投影对角线，在当前平面上的最大长度，事先计算好。
    """

    def build(self, data: Iterator, max_field=None, hierarch_r=None) -> np.ndarray:
        """
        基于层次法的识别库构建
        """
        max_field = np.inf if max_field is None else max_field
        # 计算真实三维点中心
        x, y, z = latlon2pix(
            data[0] * np.pi / 180, data[1] * np.pi / 180, self.radius, data[3]
        )
        r = data[2]
        # 计算着陆平面
        plane = np.array((x, y, z, np.ones_like(x))).T
        _, _, Vh = np.linalg.svd(plane)
        # 平面方程
        plane = Vh[-1]
        ################
        # 数值不稳定，改为直接计算
        # # 空间中以球面表达的陨石坑矩阵表示, N * 4 * 4
        # Q = np.array(
        #     [
        #         [np.ones_like(x), np.zeros_like(x), np.zeros_like(x), -x],
        #         [np.zeros_like(x), np.ones_like(x), np.zeros_like(x), -y],
        #         [np.zeros_like(x), np.zeros_like(x), np.ones_like(x), -z],
        #         [-x, -y, -z, x * x + y * y + z * z - r * r],
        #     ]
        # ).T
        # # 计算平面上的陨石坑矩阵表示，N * 3 * 3
        # C = cone_plane_cross(Q, plane)
        # 计算平面上的陨石坑的中心
        # center = center_ellipse(C)
        # center = np.array(center).T
        ################
        # 将中心的空间坐标投影为平面坐标：
        # 构建旋转矩阵，使得法向量对齐到z轴
        u = plane[:3] / np.linalg.norm(plane[:3])
        v = np.cross(u, np.array([1, 0, 0]))
        w = np.cross(u, v)
        w = w / np.linalg.norm(w)
        v = np.cross(w, u)
        R = np.column_stack((v, w, u))
        # 旋转中心
        T = -R @ np.array([0, 0, 1])
        # 旋转矩阵
        R = R.T
        # 旋转后的中心
        C = R @ np.array((x, y, z)) + T[:, None]
        center = C[:2].T
        # 计算直径的前30%分位数和后30%分位数
        if hierarch_r is None:
            hierarch_r = np.percentile(r, 30), np.percentile(r, 70)
        # 分层
        cluster_1 = np.argwhere(r <= hierarch_r[0]).squeeze()
        cluster_2 = np.argwhere((r <= hierarch_r[1]) & (r > hierarch_r[0])).squeeze()
        cluster_3 = np.argwhere(r > hierarch_r[1]).squeeze()
        ###############
        C = np.array(
            [
                [np.ones_like(C[0]), np.zeros_like(C[0]), -C[0]],
                [np.zeros_like(C[0]), np.ones_like(C[0]), -C[1]],
                [-C[0], -C[1], C[0] * C[0] + C[1] * C[1] - r * r],
            ]
        ).T
        # 分层的不变量
        invariants = []
        # 逐层计算不变量
        ## 第一层计算与自己、与第二层的不变量
        for i in cluster_1:
            ### 与自己的不变量
            # 计算当前中心的视场范围，仅取序号大于当前值的，保证不重复计算
            # 计算当前中心的视场范围，仅取序号大于当前值的，保证不重复计算
            idx = (
                np.linalg.norm(center[cluster_1] - center[i], axis=1) < max_field
            ) & (cluster_1 > i)
            # 取出当前视场范围内的椭圆参数
            # 计算不变量
            I1, I2 = self.invariant(C[i, None], C[cluster_1[idx]])
            invariants.append((np.ones_like(I1) * i, cluster_1[idx], I1, I2))
            ### 与第二层的不变量
            # 计算当前中心的视场范围，由于类别不同，所以一定不会重复计算
            idx = np.linalg.norm(center[cluster_2] - center[i], axis=1) < max_field
            if idx.sum() > 0:
                # 取出当前视场范围内的椭圆参数
                # 计算不变量
                I1, I2 = self.invariant(C[i, None], C[cluster_2[idx]])
                invariants.append((np.ones_like(I1) * i, cluster_2[idx], I1, I2))
        ## 第二层计算与自己、与第一层、与第三层的不变量
        for i in cluster_2:
            ### 与自己的不变量
            # 计算当前中心的视场范围，仅取序号大于当前值的，保证不重复计算
            idx = (
                np.linalg.norm(center[cluster_2] - center[i], axis=1) < max_field
            ) & (cluster_2 > i)
            if idx.sum() > 0:
                # 取出当前视场范围内的椭圆参数
                # 计算不变量
                I1, I2 = self.invariant(C[i, None], C[cluster_2[idx]])
                invariants.append((np.ones_like(I1) * i, cluster_2[idx], I1, I2))
            ### 与第一层的不变量已经算过了，不再重复计算
            ### 与第三层的不变量
            # 计算当前中心的视场范围，由于类别不同，所以一定不会重复计算
            idx = np.linalg.norm(center[cluster_3] - center[i], axis=1) < max_field
            if idx.sum() > 0:
                # 取出当前视场范围内的椭圆参数
                # 计算不变量
                I1, I2 = self.invariant(C[i, None], C[cluster_3[idx]])
                invariants.append((np.ones_like(I1) * i, cluster_3[idx], I1, I2))
        ## 第三层计算与自己、与第二层的不变量
        for i in cluster_3:
            ### 与自己的不变量
            # 计算当前中心的视场范围，仅取序号大于当前值的，保证不重复计算
            idx = (
                np.linalg.norm(center[cluster_3] - center[i], axis=1) < max_field
            ) & (cluster_3 > i)
            if idx.sum() > 0:
                # 取出当前视场范围内的椭圆参数
                # 计算不变量
                I1, I2 = self.invariant(C[i, None], C[cluster_3[idx]])
                invariants.append((np.ones_like(I1) * i, cluster_3[idx], I1, I2))
        invariants = np.concatenate(invariants, axis=1)
        return {
            "catalog": np.block([[x], [y], [z]]).T,
            "temp": np.block([center, r[:, None]]),
            "idx": invariants[:2].astype(np.int32),
            "descriptor": invariants[2:],
            "constants": {"max_field": max_field, "hierarch": hierarch_r},
        }

    def invariant(self, Q1, Q2):
        raise NotImplementedError
